package cn.lgwen.spark.ml.learning.kaggle;

import org.apache.spark.ml.classification.DecisionTreeClassificationModel;
import org.apache.spark.ml.classification.DecisionTreeClassifier;
import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.SparkSession;

/**
 * 2020/3/19
 * aven.wu
 * danxieai258@163.com
 * 决策树
 */
public class TitanicDecisionTreeClassify {
    public static void main(String[] args) {
        SparkSession spark = SparkSession
                .builder().master("local[*]")
                .appName("TitanicRandomForestClass")
                .getOrCreate();

        DecisionTreeClassifier dt = new DecisionTreeClassifier()
                .setLabelCol("Survived")
                .setFeaturesCol("features");
        Dataset<Row> trainSet = TitanicUtil.trainData(spark);
        DecisionTreeClassificationModel model = dt.fit(trainSet);
        Dataset<Row> testSet = TitanicUtil.testData(spark);
        Dataset<Row> res = model.transform(testSet);
        System.out.println(TitanicUtil.evaluate(res));
    }
}
